#include <ros/ros.h>
#include <sensor_msgs/Image.h>
#include <sensor_msgs/CameraInfo.h>
#include <geometry_msgs/PoseStamped.h> // Using PoseStamped for frame_id and timestamp
#include <cv_bridge/cv_bridge.h>
#include <image_transport/image_transport.h>
#include <opencv2/core.hpp>
#include <opencv2/imgproc.hpp>
#include <opencv2/calib3d.hpp>
#include <opencv2/highgui.hpp> // For imshow, waitKey (optional debug)

// For tf2 conversions
#include <tf2/LinearMath/Quaternion.h>
#include <tf2/LinearMath/Matrix3x3.h>
#include <tf2_geometry_msgs/tf2_geometry_msgs.h> // For toMsg

#include <tf2_ros/transform_broadcaster.h>  // lyt
#include <geometry_msgs/TransformStamped.h>  // lyt

class ChessboardDetector {
public:
    ChessboardDetector() : nh_(), pnh_("~"), it_(nh_), cam_info_received_(false) {
        // Get parameters
        pnh_.param("rows", pattern_rows_, 12); // Number of inner corners along height
        pnh_.param("cols", pattern_cols_, 9);   // Number of inner corners along width
        pnh_.param("marker_size", square_size_m_, 0.03);
        pnh_.param("image_is_rectified", image_is_rectified_, true);

        std::string image_topic, camera_info_topic, pose_topic_name;
        pnh_.param<std::string>("image_topic", image_topic, "/image");
        pnh_.param<std::string>("camera_info_topic", camera_info_topic, "/camera_info");
        pnh_.param<std::string>("pose_topic", pose_topic_name, "/chessboard_pose");
        pnh_.param<bool>("show_debug_image", show_debug_image_, true);

        pattern_size_ = cv::Size(pattern_cols_, pattern_rows_);

        // Prepare object points (3D points of chessboard corners in its own coordinate system)
        // (0,0,0), (square_size,0,0), ..., ( (cols-1)*square_size, (rows-1)*square_size, 0 )
        for (int i = 0; i < pattern_rows_; ++i) {
            for (int j = 0; j < pattern_cols_; ++j) {
                object_points_.push_back(cv::Point3f(j * square_size_m_, i * square_size_m_, 0.0f));
            }
        }

        // Subscribers
        // Use image_transport for robust image subscription
        image_sub_ = it_.subscribe(image_topic, 1, &ChessboardDetector::imageCallback, this);
        cam_info_sub_ = nh_.subscribe(camera_info_topic, 1, &ChessboardDetector::cameraInfoCallback, this);

        // Publisher
        pose_pub_ = nh_.advertise<geometry_msgs::PoseStamped>(pose_topic_name, 10);
        if (show_debug_image_) {
            debug_image_pub_ = it_.advertise("/chessboard_detector/debug_image", 1);
            cv::namedWindow("Chessboard Detection", cv::WINDOW_AUTOSIZE);
        }

        ROS_INFO("Chessboard detector initialized.");
        ROS_INFO("Pattern: %d x %d, Square size: %.3f m", pattern_cols_, pattern_rows_, square_size_m_);
        ROS_INFO("Subscribing to image: %s", image_sub_.getTopic().c_str());
        ROS_INFO("Subscribing to camera info: %s", cam_info_sub_.getTopic().c_str());
        ROS_INFO("Publishing pose to: %s", pose_pub_.getTopic().c_str());
    }

    ~ChessboardDetector() {
        if (show_debug_image_) {
            cv::destroyWindow("Chessboard Detection");
        }
    }

private:
    void cameraInfoCallback(const sensor_msgs::CameraInfoConstPtr& msg) {
        if (cam_info_received_) {
            // Optionally, you might only want to process this once if camera info is static
            // cam_info_sub_.shutdown(); 
            return;
        }

        camera_matrix_ = cv::Mat(3, 3, CV_64F, (void*)msg->K.data()).clone();
        if (image_is_rectified_) {
            // If image is rectified, distortion coefficients are typically all zeros or not used.
            // RealSense color images are usually rectified.
            dist_coeffs_ = cv::Mat::zeros(1, 5, CV_64F); 
        } else {
            dist_coeffs_ = cv::Mat(msg->D.size(), 1, CV_64F, (void*)msg->D.data()).clone();
        }

        cam_info_received_ = true;
        ROS_INFO("Camera info received.");
        // It's good practice to shut down the subscriber if info is static
        // cam_info_sub_.shutdown(); 
        // ROS_INFO("Camera info subscriber shut down as info is static.");
    }

    void imageCallback(const sensor_msgs::ImageConstPtr& msg) {
        if (!cam_info_received_) {
            ROS_WARN_THROTTLE(5.0, "Waiting for camera info...");
            return;
        }

        cv_bridge::CvImagePtr cv_ptr;
        try {
            cv_ptr = cv_bridge::toCvCopy(msg, sensor_msgs::image_encodings::BGR8);
        } catch (cv_bridge::Exception& e) {
            ROS_ERROR("cv_bridge exception: %s", e.what());
            return;
        }

        cv::Mat gray_image;
        cv::cvtColor(cv_ptr->image, gray_image, cv::COLOR_BGR2GRAY);

        std::vector<cv::Point2f> image_points;
        bool found = cv::findChessboardCorners(gray_image, pattern_size_, image_points,
                                               cv::CALIB_CB_ADAPTIVE_THRESH |
                                               cv::CALIB_CB_NORMALIZE_IMAGE |
                                               cv::CALIB_CB_FAST_CHECK); // FAST_CHECK speeds up if no board

        if (found) {
            // Refine corner locations
            cv::cornerSubPix(gray_image, image_points, cv::Size(11, 11), cv::Size(-1, -1),
                             cv::TermCriteria(cv::TermCriteria::EPS + cv::TermCriteria::COUNT, 30, 0.1));

            cv::Mat rvec, tvec;
            // Solve for pose
            // Note: object_points_ are in chessboard frame, image_points are in image pixels
            // solvePnP returns rvec and tvec that transform points from model (object) frame to camera frame
            bool solved = cv::solvePnP(object_points_, image_points, camera_matrix_, dist_coeffs_, rvec, tvec);

            if (solved) {
                geometry_msgs::PoseStamped pose_stamped_msg;
                pose_stamped_msg.header.stamp = msg->header.stamp; // Use image timestamp
                pose_stamped_msg.header.frame_id = msg->header.frame_id; // Camera optical frame

                // Translation vector (tvec) directly maps to position
                pose_stamped_msg.pose.position.x = tvec.at<double>(0);
                pose_stamped_msg.pose.position.y = tvec.at<double>(1);
                pose_stamped_msg.pose.position.z = tvec.at<double>(2);

                // Rotation vector (rvec) needs to be converted to quaternion
                cv::Mat R_matrix;
                cv::Rodrigues(rvec, R_matrix); // Convert rotation vector to rotation matrix

                tf2::Matrix3x3 tf_R_matrix(
                    R_matrix.at<double>(0,0), R_matrix.at<double>(0,1), R_matrix.at<double>(0,2),
                    R_matrix.at<double>(1,0), R_matrix.at<double>(1,1), R_matrix.at<double>(1,2),
                    R_matrix.at<double>(2,0), R_matrix.at<double>(2,1), R_matrix.at<double>(2,2)
                );

                tf2::Quaternion tf_q;
                tf_R_matrix.getRotation(tf_q);
                tf_q.normalize(); // Ensure it's a unit quaternion

                pose_stamped_msg.pose.orientation = tf2::toMsg(tf_q);

                pose_pub_.publish(pose_stamped_msg);

                // 创建 TransformStamped 消息
                geometry_msgs::TransformStamped transformStamped;
                transformStamped.header.stamp = msg->header.stamp;
                transformStamped.header.frame_id = msg->header.frame_id; // 相机光学帧
                transformStamped.child_frame_id = "chessboard_frame";  // 棋盘坐标系

                transformStamped.transform.translation.x = tvec.at<double>(0);
                transformStamped.transform.translation.y = tvec.at<double>(1);
                transformStamped.transform.translation.z = tvec.at<double>(2);
                transformStamped.transform.rotation = tf2::toMsg(tf_q);

                // 在 imageCallback 函数中，发布 TF 变换之前添加日志信息
                // ROS_INFO("Publishing TF transform from %s to chessboard_frame", msg->header.frame_id.c_str());
                // tf_broadcaster_.sendTransform(transformStamped);

                // 发布TF变换
                tf_broadcaster_.sendTransform(transformStamped);

            } else {
                ROS_WARN_THROTTLE(1.0, "solvePnP failed.");
            }

            if (show_debug_image_ || debug_image_pub_.getNumSubscribers() > 0) {
                 cv::drawChessboardCorners(cv_ptr->image, pattern_size_, image_points, found);
            }

        } else {
            ROS_WARN_THROTTLE(5.0, "Chessboard not found in image.");
        }

        if (show_debug_image_) {
            cv::imshow("Chessboard Detection", cv_ptr->image);
            cv::waitKey(1);
        }
        if (debug_image_pub_.getNumSubscribers() > 0) {
            debug_image_pub_.publish(cv_ptr->toImageMsg());
        }
    }

    ros::NodeHandle nh_;
    ros::NodeHandle pnh_; // Private NodeHandle for parameters
    image_transport::ImageTransport it_;

    image_transport::Subscriber image_sub_;
    ros::Subscriber cam_info_sub_;
    ros::Publisher pose_pub_;
    image_transport::Publisher debug_image_pub_;


    cv::Size pattern_size_;
    int pattern_rows_;
    int pattern_cols_;
    double square_size_m_;
    bool image_is_rectified_;
    bool show_debug_image_;

    std::vector<cv::Point3f> object_points_;
    cv::Mat camera_matrix_;
    cv::Mat dist_coeffs_;
    bool cam_info_received_;

    tf2_ros::TransformBroadcaster tf_broadcaster_;
};

int main(int argc, char** argv) {
    // ros::init(argc, argv, "chessboard_detector_node");
    ros::init(argc, argv, "chessboard_detector");
    ChessboardDetector detector;
    ros::spin();
    return 0;
}